package com.treasure.infrastructure.repository.extraction.vector;

import java.time.LocalDateTime;

import ch.qos.logback.classic.Level;
import ch.qos.logback.classic.spi.Configurator;
import com.alibaba.dashscope.embeddings.TextEmbeddingResultItem;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.google.common.collect.Lists;
import com.treasure.common.utils.TextCleaner;
import com.treasure.infrastructure.mapstruct.KnowledgeDocumentAssembler;
import com.treasure.common.model.UploadDocumentRequest;
import com.treasure.infrastructure.model.doc.SegmentVectorDocBean;
import com.treasure.infrastructure.domain.llm.qianwen.client.QwenEmbeddingVectorClient;
import com.treasure.infrastructure.model.po.KnowledgeDocumentVector;
import com.treasure.common.utils.EmbeddingConvert;
import com.treasure.infrastructure.repository.dao.mapper.KnowledgeDocumentVectorMapper;
import lombok.RequiredArgsConstructor;
import org.apache.logging.log4j.LogManager;
import org.elasticsearch.common.Strings;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @author sunhongmin
 * @date 2023/8/25 17:50
 * @description
 */
@Service
@RequiredArgsConstructor
public class SegmentVectorDocRepositoryImpl implements SegmentVectorDocRepository {
    private final QwenEmbeddingVectorClient embeddingVectorClient;
    private final KnowledgeDocumentVectorMapper knowledgeDocumentVectorMapper;


    @Override
    public void indexVector(List<SegmentVectorDocBean> docSegmentList, UploadDocumentRequest uploadDocumentRequest) throws NoApiKeyException {
        if (CollectionUtils.isEmpty(docSegmentList)) {
            return;
        }

        List<String> passagesList = docSegmentList.stream().map(SegmentVectorDocBean::getPassages).collect(Collectors.toList());
        List<List<String>> passagesPartition = Lists.partition(passagesList, 20);
        for (List<String> passagesPartList : passagesPartition) {
            List<TextEmbeddingResultItem> embeddingResultItems = embeddingVectorClient.textEmbeddingVector(passagesPartList);
            if (CollectionUtils.isEmpty(embeddingResultItems)) {
                continue;
            }
            List<KnowledgeDocumentVector> insertList = new ArrayList<>(embeddingResultItems.size());
            for (int i = 0; i < embeddingResultItems.size(); i++) {
                TextEmbeddingResultItem embeddingResultItem = embeddingResultItems.get(i);
                KnowledgeDocumentVector knowledgeDocumentVector = getKnowledgeDocumentVector(uploadDocumentRequest, passagesPartList.get(i), embeddingResultItem);
                insertList.add(knowledgeDocumentVector);
            }
            if (!CollectionUtils.isEmpty(insertList)) {
                knowledgeDocumentVectorMapper.batchInsertDoc(insertList);
            }
        }
    }

    @NotNull
    private static KnowledgeDocumentVector getKnowledgeDocumentVector(UploadDocumentRequest uploadDocumentRequest, String passageContent, TextEmbeddingResultItem embeddingResultItem) {
        KnowledgeDocumentVector knowledgeDocumentVector = new KnowledgeDocumentVector();
        knowledgeDocumentVector.setCorpId(uploadDocumentRequest.getCorpId());
        knowledgeDocumentVector.setDocumentId(uploadDocumentRequest.getDocumentId().intValue());
        knowledgeDocumentVector.setDocumentName(uploadDocumentRequest.getDocumentName());
        knowledgeDocumentVector.setPassages(TextCleaner.cleanText(passageContent));
        knowledgeDocumentVector.setEmbedding(EmbeddingConvert.toString(embeddingResultItem.getEmbedding()));
        knowledgeDocumentVector.setCreateTime(LocalDateTime.now());
        knowledgeDocumentVector.setHasDeleted(false);
        knowledgeDocumentVector.setCreatorId(uploadDocumentRequest.getUserId());
        return knowledgeDocumentVector;
    }


    @Override
    public List<SegmentVectorDocBean> searchByVector(String corpId, String question) throws NoApiKeyException {
        if (!Strings.hasLength(question)) {
            return Collections.emptyList();
        }
        List<TextEmbeddingResultItem> embeddingResultItems = embeddingVectorClient.textEmbeddingVector(Collections.singletonList(question));
        TextEmbeddingResultItem firstVector = embeddingResultItems.stream().findFirst().orElse(null);
        if (Objects.isNull(firstVector)) {
            return Collections.emptyList();
        }
        List<KnowledgeDocumentVector> knowledgeDocumentVectors = knowledgeDocumentVectorMapper.searchByVector(corpId, EmbeddingConvert.toString(firstVector.getEmbedding()));
        return KnowledgeDocumentAssembler.INSTANCE.convertToSegmentVectorDocBeanList(knowledgeDocumentVectors);
    }


}
